Fix dynamic wires with samples #1321

rniczh · 2026-01-06T20:51:45Z

Context:

from catalyst import *
import pennylane as qml

qml.capture.enable()

@qjit
@qml.qnode(qml.device("lightning.qubit", wires=3, shots=10))
def circuit():
    with qml.allocate(2) as qs:
        qml.X(qs[1])

    return qml.sample(wires=[0, 1])

print(circuit())

It get unexpected result due to the aux wires involved.

[[0 0]
 [0 1]
 [0 0]
 [1 0]
 [0 0]
 [0 0]
 [0 1]
 [0 0]
 [1 0]
 [0 0]]

Description of the Change:

~~Only take the device wires for generating samples~~
Compact state vector before measurement.

For state vector normalization:

$|\psi| = \sqrt{\sum_i |\psi_i|^2}$
$normalized_i = \psi_i / |\psi|$

Benefits:

Possible Drawbacks:

Related GitHub Issues:

PennyLaneAI/catalyst#2339
[[sc-107206]]

github-actions · 2026-01-06T20:51:57Z

Hello. You may have forgotten to update the changelog!
Please edit .github/CHANGELOG.md with:

A one-to-two sentence description of the change. You may include a small working example for new features.
A link back to this PR.
Your name (or GitHub username) in the contributors section.

dime10

Thank you @rniczh !

pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.cpp

dime10 · 2026-01-06T21:24:40Z

pennylane_lightning/core/simulators/lightning_qubit/measurements/MeasurementsLQubit.hpp

+        const std::vector<std::size_t> &wires, const std::string &kernelname,
+        std::size_t num_burnin, std::size_t num_samples) {
        std::uniform_real_distribution<PrecisionT> distrib(0.0, 1.0);
+        std::size_t num_qubits = wires.size();


Do we know that this is sufficient to identify which wires should be selected in a more complex dynamic (de)allocation scheme, or does it rely on the assumption that the selected wires are always the bottom n wires?

Mentioned this offline but I think we'd be better off providing a universal solution that we never got around to implementing (see #1254 (comment)).

I updated the PR, basically removed the original solution that only pass the device_wires to GenerateSample. Now, it will compact the state vector before the measurement process to ensure the erased qubits are erased from state vector as well.
~~ca0ff43~~ Edit: 77fbbe9

Wow that was fast! 🤩
@maliasadi are you able to look this over?

codecov · 2026-01-07T20:06:30Z

Codecov Report

❌ Patch coverage is 99.58333% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 93.89%. Comparing base (e7d82a2) to head (6c8cdb2).

Files with missing lines	Patch %	Lines
...rs/lightning_qubit/catalyst/LightningSimulator.cpp	98.00%	1 Missing ⚠️

Additional details and impacted files

@@              Coverage Diff               @@
##           v0.44.0_rc    #1321      +/-   ##
==============================================
- Coverage       96.01%   93.89%   -2.13%     
==============================================
  Files             307      243      -64     
  Lines           46923    40264    -6659     
==============================================
- Hits            45055    37804    -7251     
- Misses           1868     2460     +592

Flag	Coverage Δ
unit_tests	`93.89% <99.58%> (-2.13%)`	⬇️

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

maliasadi

Thanks @rniczh! A few things:

If we need this to go out with the release, you will need to rebase it against #1317 and sync ASAP with @jzaia18
Update Lighnting-Kokkos and Lightning-GPU catalyst impl similar to Lightning-Qubit
Please add a changelog entry for this fix

--------- Co-authored-by: Ali Asadi <[email protected]>

jzaia18

Nice work Hongsheng! Just a few questions and comments. ~~Is there any reason we're implementing this for lightning.qubit but not the other devices? Does this issue exist for the others as well?~~ nvm, just saw Ali already addressed this

pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.cpp

pennylane_lightning/core/catalyst/LightningQubitManager.hpp

pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.cpp

.github/CHANGELOG.md

pennylane_lightning/core/simulators/lightning_gpu/catalyst/LightningGPUSimulator.cpp

maliasadi · 2026-01-08T22:23:19Z

pennylane_lightning/core/simulators/lightning_gpu/catalyst/LightningGPUSimulator.cpp

+        *(this->device_sv)};
+}
+
+void LightningGPUSimulator::CompactStateVector() {


Where are these tested? I couldn't find any tests for these methods for all devices.

The function is triggered by getMeasurements(), ensuring that any test executing the measurement process also updates the state vector. Currently, only one test case is added (the one that reported from the issue PennyLaneAI/catalyst#2339). The specific test can be found here:

https://github.com/PennyLaneAI/pennylane-lightning/pull/1321/changes#diff-6f0e99f87b0a9753689936392636b52eb19b81aac144b24d12de297f3c617447R335

Will this be added to Catalyst for the release?

per our offline discussions, what if these aux wires are entangled, we should test and add some handling mechanism for these cases in both Lightning & Catalyst?

pyproject.toml

pennylane_lightning/core/catalyst/LightningQubitManager.hpp

pennylane_lightning/core/simulators/lightning_gpu/catalyst/LightningGPUSimulator.cpp

pennylane_lightning/core/catalyst/LightningQubitManager.hpp

maliasadi · 2026-01-08T22:39:47Z

pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.hpp

    auto GenerateSamples(size_t shots) -> std::vector<size_t>;

+    // Compact state vector by removing released qubits
+    void CompactStateVector();


minor: this is the style we've been following for a while:

Suggested change

void CompactStateVector();

void compactStateVector();

Should we call this reducedStateVector instead of compact?

Sure! I updated the name here: a5e2132

jzaia18 · 2026-01-08T23:08:25Z

Could you also bump the rc version to 3 within pennylane_lightning/core/_version.py?

Co-authored-by: Ali Asadi <[email protected]>

rniczh · 2026-01-09T01:13:27Z

Could you also bump the rc version to 3 within pennylane_lightning/core/_version.py?

Sure abd373d

jzaia18

Looking good, just 1 last comment

pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.cpp

This comment has been minimized.

Sign in to view

dime10 reviewed Jan 6, 2026

View reviewed changes

This comment has been minimized.

Sign in to view

maliasadi self-requested a review January 8, 2026 15:06

maliasadi reviewed Jan 8, 2026

View reviewed changes

rniczh changed the base branch from master to v0.44.0_rc January 8, 2026 15:19

rniczh added 5 commits January 8, 2026 10:24

rebase

79451be

Fix dynamic wires with samples

7078d0e

fix

ddfdb53

add back the comment

9d19fa9

Compact state vector before measurement

7ae4e5e

jzaia18 added this to the 0.44.0 milestone Jan 8, 2026

jzaia18 and others added 8 commits January 8, 2026 10:28

Rebase

b3e806f

--------- Co-authored-by: Ali Asadi <[email protected]>

update

d365d2d

fix

6771882

no need to collapse released qubit during the

0e11617

fix compact timing

635370d

should use big-endian

aa98ec9

Add test

a8367bd

update test

6e9f177

rniczh force-pushed the rniczh/fix-dynamic-wires-with-samples branch from d3abf70 to 6e9f177 Compare January 8, 2026 15:32

rniczh changed the base branch from v0.44.0_rc to master January 8, 2026 15:37

rniczh changed the base branch from master to v0.44.0_rc January 8, 2026 15:37

clean

5bac5ae

jzaia18 reviewed Jan 8, 2026

View reviewed changes

rniczh added 2 commits January 8, 2026 13:32

update

54a87e1

fix

14a351e

rniczh added 2 commits January 8, 2026 14:56

revert

90d9300

fix coverage

bf6ff01

rniczh removed the ci:use-gpu-runner Enable usage of GPU runner for this Pull Request label Jan 8, 2026

maliasadi reviewed Jan 8, 2026

View reviewed changes

jzaia18 reviewed Jan 8, 2026

View reviewed changes

pennylane_lightning/core/catalyst/LightningQubitManager.hpp Outdated Show resolved Hide resolved

maliasadi reviewed Jan 8, 2026

View reviewed changes

pennylane_lightning/core/simulators/lightning_gpu/catalyst/LightningGPUSimulator.cpp Outdated Show resolved Hide resolved

maliasadi reviewed Jan 8, 2026

View reviewed changes

rniczh and others added 8 commits January 8, 2026 19:27

Update .github/CHANGELOG.md

203a53b

Co-authored-by: Ali Asadi <[email protected]>

Update .github/CHANGELOG.md

1ed47be

Co-authored-by: Ali Asadi <[email protected]>

remove redundant

81c4859

formatting

6eb8caa

compactStateVector -> reducedStateVector

a5e2132

Update pennylane_lightning/core/catalyst/LightningQubitManager.hpp

43a1b6a

Co-authored-by: Ali Asadi <[email protected]>

formatting

9fc0cc0

bumping to rc3

abd373d

rniczh added 3 commits January 8, 2026 21:31

move normalizeStateVector to Util

60511a8

fix util

128dfff

no need to cover the complex with custom allocator

57b3a61

rniczh mentioned this pull request Jan 9, 2026

Raise error message for dynamic wires with non-prob measurement PennyLaneAI/catalyst#2358

Merged

revert project.toml

0ec104f

rniczh added the ci:use-gpu-runner Enable usage of GPU runner for this Pull Request label Jan 9, 2026

rniczh added 2 commits January 9, 2026 00:27

ci

fcc5229

fix changelog

2b784c7

jzaia18 reviewed Jan 9, 2026

View reviewed changes

pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.cpp Outdated Show resolved Hide resolved

rniczh added 2 commits January 9, 2026 11:35

change fn name of reducedStateVector

0f53498

change to use fn name reduceStateVector for all devices

6c8cdb2

paul0403 requested a review from mlxd January 9, 2026 18:43

jzaia18 removed this from the 0.44.0 milestone Jan 9, 2026

Fix dynamic wires with samples #1321

Are you sure you want to change the base?

Fix dynamic wires with samples #1321

Uh oh!

Conversation

rniczh commented Jan 6, 2026 • edited by sengthai Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Jan 6, 2026

Uh oh!

This comment has been minimized.

dime10 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rniczh Jan 7, 2026 • edited by dime10 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This comment has been minimized.

codecov bot commented Jan 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

This comment has been minimized.

maliasadi left a comment

Choose a reason for hiding this comment

Uh oh!

jzaia18 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

maliasadi Jan 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jzaia18 commented Jan 8, 2026

Uh oh!

rniczh commented Jan 9, 2026

Uh oh!

jzaia18 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

rniczh commented Jan 6, 2026 •

edited by sengthai

Loading

rniczh Jan 7, 2026 •

edited by dime10

Loading

codecov bot commented Jan 7, 2026 •

edited

Loading

jzaia18 left a comment •

edited

Loading

maliasadi Jan 9, 2026 •

edited

Loading